﻿/*
 * @author: wizardc
 */

using System.IO;
using Dou.GPU;
using Dou.Job;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Jobs;
using Unity.Mathematics;
using UnityEngine;

namespace Demo.WorldMap.GPU
{
    /// <summary>
    /// GPU 类型渲染类，渲染一颗 BVH 树
    /// 注：每个 SceneStaticGPURender 相互之间都是独立的，即 2 个 SceneStaticGPURender 之间不会有合批的情况发生
    /// </summary>
    public unsafe class SceneStaticGPURender : ISceneStaticRender
    {
        public GPUStaticRenderTree tree;
        public Stream dataStream;
        
        public SceneStaticCameraCtl cameraCtl;
        
        private GPUDrawContainer _container;
        
        private NativeArray<bool> _visibleFlagArray; // 节点显示状态数组
        private NativeList<int> _visibleNodeList; // 显示节点索引列表
        private NativeList<int> _addNodeList; // 新增显示的节点索引列表
        private NativeList<int> _removeNodeList; // 移除显示的节点索引列表
        
        private NodeVisibleUpdateJob _nodeManagedJob;
        private ManagedJob<NodeVisibleUpdateJob> _nodeJob;
        private JobHandle _nodeJobHandle;
        
        private DataReductionJob _dataManagedJob;
        private ManagedJob<DataReductionJob> _dataJob;
        private JobHandle _dataJobHandle;

        private RemoveDataSortJob _sortManagedJob;
        private ManagedJobParallelFor<RemoveDataSortJob> _sortJob;
        private JobHandle _sortJobHandle;

        private FillBatchDataJob _fillManagedJob;
        private ManagedJob<FillBatchDataJob> _fillJob;
        private JobHandle _fillJobHandle;
        
        private ulong _renderNodeHandler;
        private ulong _renderLeafHandler;

        private bool _needLateUpdate = false;
        
        private bool _enable = true;
        private int _lod = 0;
        private bool _forceUpdate = false;
        
        public void SetEnable(bool enable)
        {
            _enable = enable;
            if (_container != null && !_enable)
            {
                _container.ClearBatchData();
            }
        }
        
        public void SetLOD(int lod)
        {
            _lod = lod;
            _container?.SetLOD(_lod);
        }
        
        public void SetForceUpdate()
        {
            _forceUpdate = true;
        }
        
        public void OnEnable()
        {
            _container = new GPUDrawContainer();
            _container.Init(tree.res, tree.lodRes, dataStream);
            _container.SetLOD(_lod);
            
            _visibleFlagArray = new NativeArray<bool>(tree.nodes.Length, Allocator.Persistent);
            _visibleNodeList = new NativeList<int>(Allocator.Persistent);
            _addNodeList = new NativeList<int>(Allocator.Persistent);
            _removeNodeList = new NativeList<int>(Allocator.Persistent);
            
            _nodeManagedJob = new NodeVisibleUpdateJob
            {
                cameraCtl = cameraCtl,
                visibleFlagArray = _visibleFlagArray,
                addNodeList = _addNodeList,
                removeNodeList = _removeNodeList
            };
            _nodeJob = new ManagedJob<NodeVisibleUpdateJob>();
            _nodeJob.Work = _nodeManagedJob;
            
            _dataManagedJob = new DataReductionJob
            {
                _container = _container,
                visibleNodeList = _visibleNodeList,
                addNodeList = _addNodeList,
                removeNodeList = _removeNodeList
            };
            _dataJob = new ManagedJob<DataReductionJob>();
            _dataJob.Work = _dataManagedJob;
            
            _sortManagedJob = new RemoveDataSortJob
            {
                _container = _container,
                visibleNodeList = _visibleNodeList
            };
            _sortJob = new ManagedJobParallelFor<RemoveDataSortJob>();
            _sortJob.Work = _sortManagedJob;
            
            _fillManagedJob = new FillBatchDataJob
            {
                cameraCtl = cameraCtl,
                _container = _container
            };
            _fillJob = new ManagedJob<FillBatchDataJob>();
            _fillJob.Work = _fillManagedJob;
        }

        public void Update()
        {
            if (!_enable)
            {
                return;
            }
            _needLateUpdate = true;
            
            var nodePtr = (GPUStaticRenderNode*) UnsafeUtility.PinGCArrayAndGetDataAddress(tree.nodes, out _renderNodeHandler);
            var leafPtr = (GPUStaticRenderLeaf*) UnsafeUtility.PinGCArrayAndGetDataAddress(tree.leaves, out _renderLeafHandler);
            
            _nodeManagedJob.nodePtr = nodePtr;
            _nodeManagedJob.forceUpdate = _forceUpdate;
            _nodeJobHandle = _nodeJob.Schedule();
            _nodeJobHandle.Complete();

            _dataManagedJob.nodePtr = nodePtr;
            _dataManagedJob.leafPtr = leafPtr;
            _dataJobHandle = _dataJob.Schedule();
            _dataJobHandle.Complete();
            
            if (_dataManagedJob.isDirty)
            {
                _sortManagedJob.nodePtr = nodePtr;
                _sortManagedJob.leafPtr = leafPtr;
                _sortJobHandle = _sortJob.Schedule(_container.resCount, 1);
                JobHandle.ScheduleBatchedJobs();
            }
        }

        public void LateUpdate()
        {
            if (!_needLateUpdate)
            {
                return;
            }
            _needLateUpdate = false;
            
            if (_dataManagedJob.isDirty)
            {
                _sortJobHandle.Complete();
            }
            
            UnsafeUtility.ReleaseGCObject(_renderNodeHandler);
            UnsafeUtility.ReleaseGCObject(_renderLeafHandler);
            
            // 删除不需要的数据，同时从数据流中加载需要的数据
            _container.CopyData();
            // 清除所有绘制批次的数据，这里的优化点是可以通过标脏来处理仅发生变化的数据
            _container.ClearBatchData();
            
            // 填充所有绘制批次的数据
            _fillJobHandle = _fillJob.Schedule();
            _fillJobHandle.Complete();
            
            // 资源加载必须在主线程执行
            _container.TryLoadRes();
            
            _forceUpdate = false;
        }
        
        public void OnDisable()
        {
            _visibleFlagArray.Dispose();
            _visibleNodeList.Dispose();
            _addNodeList.Dispose();
            _removeNodeList.Dispose();
            _nodeManagedJob = null;
            _nodeJob.Dispose();
            _dataManagedJob = null;
            _dataJob.Dispose();
            _sortManagedJob = null;
            _sortJob.Dispose();
            _fillManagedJob = null;
            _fillJob.Dispose();
            _container.Dispose();
        }

#if UNITY_EDITOR

        public void DrawObjAABB()
        {
            // 绘制没有视椎裁剪的所有载入数据
            Gizmos.color = Color.yellow;
            {
                var count = _container.resCount;
                for (int i = 0; i < count; i++)
                {
                    var resInfo = _container.GetResInfo(i);
                    var ptr = (float4x4*) resInfo.transData.GetUnsafePtr();
                    var num = resInfo.transData.Length / GPUDefine.Stride_Float4x4;
                    for (int j = 0; j < num; j++)
                    {
                        var p = *(ptr + j);
                        Gizmos.DrawWireCube(new Vector3(p.c3.x, p.c3.y, p.c3.z), new Vector3(p.c0.w * 2 + 0.1f, p.c1.w * 2 + 0.1f, p.c0.w * 2 + 0.1f)); // 略大一点，和视椎裁剪的数据错开
                    }
                }
            }
            
            // 绘制最终会呈现到场景的数据
            Gizmos.color = Color.green;
            {
                var count = _container.batchCount;
                for (int i = 0; i < count; i++)
                {
                    var batchDraw = _container.GetBatchDraw(i);
                    batchDraw.DrawObjAABB();
                }
            }
        }
        
#endif
        
    }
}
